#-
# ==========================================================================
# Copyright (C) 1995 - 2006 Autodesk, Inc. and/or its licensors.  All 
# rights reserved.
#
# The coded instructions, statements, computer programs, and/or related 
# material (collectively the "Data") in these files contain unpublished 
# information proprietary to Autodesk, Inc. ("Autodesk") and/or its 
# licensors, which is protected by U.S. and Canadian federal copyright 
# law and by international treaties.
#
# The Data is provided for use exclusively by You. You have the right 
# to use, modify, and incorporate this Data into other products for 
# purposes authorized by the Autodesk software license agreement, 
# without fee.
#
# The copyright notices in the Software and this entire statement, 
# including the above license grant, this restriction and the 
# following disclaimer, must be included in all copies of the 
# Software, in whole or in part, and all derivative works of 
# the Software, unless such copies or derivative works are solely 
# in the form of machine-executable object code generated by a 
# source language processor.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. 
# AUTODESK DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED 
# WARRANTIES INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF 
# NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR 
# PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE, OR 
# TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS LICENSORS 
# BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL, 
# DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK 
# AND/OR ITS LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY 
# OR PROBABILITY OF SUCH DAMAGES.
#
# ==========================================================================
#+

###############################################################################
##
## basicShape.py
##
## Description:
##    Registers a new type of shape with maya called "basicShape".
##    This shape will display rectangles, triangles, and circles using basic gl
##
##
##    There are no output attributes for this shape.
##    The following input attributes define the type of shape to draw.
##
##       shapeType  : 0=rectangle, 1=circle, 2=triangle
##       radius		: circle radius
##       height		: rectangle and triangle height
##		 width		: rectangle and triangle width
##
################################################################################

# Usage:
# import maya
# maya.cmds.loadPlugin("basicShape.py")
# maya.cmds.createNode("spBasicShape")
#
# An object will be created with reference to the node.  By default it will be a rectangle.
# Use the different options node options to change the type of shape or its size and shape.
# Add textures or manipulate as with any other object.
#

import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaRender as OpenMayaRender
import maya.OpenMayaUI as OpenMayaUI

import math
import sys


kPluginNodeTypeName = "spBasicShape"
spBasicShapeNodeId = OpenMaya.MTypeId(0x87018)

glRenderer = OpenMayaRender.MHardwareRenderer.theRenderer()
glFT = glRenderer.glFunctionTable()

kLeadColor 				= 18 # green
kActiveColor			= 15 # white
kActiveAffectedColor	= 8  # purple
kDormantColor			= 4  # blue
kHiliteColor			= 17 # pale blue

kDefaultRadius = 1.0
kDefaultHeight = 2.0
kDefaultWidth = 2.0
kDefaultShapeType = 0


#####################################################################
##
## Geometry class
##
class basicGeom:
	radius = kDefaultRadius
	height = kDefaultHeight
	width = kDefaultWidth
	shapeType = kDefaultShapeType


#####################################################################
##
## Shape class - defines the non-UI part of a shape node
##
class basicShape(OpenMayaMPx.MPxSurfaceShape):
	def __init__(self):
		OpenMayaMPx.MPxSurfaceShape.__init__(self)

		# class variables
		aShapeType = OpenMaya.MObject()
		aRadius = OpenMaya.MObject()
		aHeight = OpenMaya.MObject()
		aWidth = OpenMaya.MObject()

		# geometry
		self.__myGeometry = basicGeom()


	# override
	def postConstructor(self):
		"""
		 When instances of this node are created internally, the MObject associated
		 with the instance is not created until after the constructor of this class
		 is called. This means that no member functions of MPxSurfaceShape can
		 be called in the constructor.
		 The postConstructor solves this problem. Maya will call this function
		 after the internal object has been created.
		 As a general rule do all of your initialization in the postConstructor.
		"""
		self.setRenderable(True)


	# override
	def compute(self, plug, dataBlock):
		"""
		 Since there are no output attributes this is not necessary but
		 if we wanted to compute an output mesh for rendering it would
		 be done here base on the inputs.
		"""
		return OpenMaya.kUnknownParameter


	# override
	def getInternalValue(self, plug, datahandle):
		"""
		 Handle internal attributes.
		 In order to impose limits on our attribute values we
		 mark them internal and use the values in fGeometry intead.
		"""
		if (plug == basicShape.aRadius):
			datahandle.setDouble(self.__myGeometry.radius)

		elif (plug == basicShape.aHeight):
			datahandle.setDouble(self.__myGeometry.height)
		
		elif (plug == basicShape.aWidth):
			datahandle.setDouble(self.__myGeometry.width)

		else:
			return OpenMayaMPx.MPxSurfaceShape.getInternalValue(self, plug, datahandle)

		return True


	# override
	def setInternalValue(self, plug, datahandle):
		"""
		 Handle internal attributes.
		 In order to impose limits on our attribute values we
		 mark them internal and use the values in fGeometry intead.
		"""

		# the minimum radius is 0
		#
		if (plug == basicShape.aRadius):
			radius = datahandle.asDouble()
			
			if (radius < 0):
				radius = 0

			self.__myGeometry.radius = radius

		elif (plug == basicShape.aHeight):
			val = datahandle.asDouble()
			if (val <= 0):
				val = 0.1
			self.__myGeometry.height = val

		elif (plug == basicShape.aWidth):
			val = datahandle.asDouble()
			if (val <= 0):
				val = 0.1
			self.__myGeometry.width = val
		
		else:
			return OpenMayaMPx.MPxSurfaceShape.setInternalValue(self, plug, datahandle)
		
		return True


	# override
	def isBounded(self):
		return True


	# override
	def boundingBox(self):
		"""
		 Returns the bounding box for the shape.
		 In this case just use the radius and height attributes
		 to determine the bounding box.
		"""
		result = OpenMaya.MBoundingBox()

		geom = self.geometry()

		r = geom.radius
		result.expand(OpenMaya.MPoint(r,r,r))
		result.expand(OpenMaya.MPoint(-r,-r,-r))
		
		r = geom.height/2.0
		result.expand(OpenMaya.MPoint(r,r,r))
		result.expand(OpenMaya.MPoint(-r,-r,-r))
		
		r = geom.width/2.0
		result.expand(OpenMaya.MPoint(r,r,r))
		result.expand(OpenMaya.MPoint(-r,-r,-r))

		return result


	def geometry(self):
		"""
		 This function gets the values of all the attributes and
		 assigns them to the fGeometry. Calling MPlug::getValue
		 will ensure that the values are up-to-date.
		"""
		# return self.__myGeometry

		this_object = self.thisMObject()

		plug = OpenMaya.MPlug(this_object, basicShape.aRadius)
		self.__myGeometry.radius = plug.asDouble()

		plug.setAttribute(basicShape.aHeight)
		self.__myGeometry.height = plug.asDouble()
		
		plug.setAttribute(basicShape.aWidth)
		self.__myGeometry.width = plug.asDouble()
		
		plug.setAttribute(basicShape.aShapeType)
		self.__myGeometry.shapeType = plug.asShort() # enum????

		return self.__myGeometry

def printMsg(msg):
	print msg
	stream=OpenMaya.MStreamUtils.stdOutStream()
	OpenMaya.MStreamUtils.writeCharBuffer(stream,msg)

#####################################################################
##
## UI class	- defines the UI part of a shape node
##
class basicShapeUI(OpenMayaMPx.MPxSurfaceShapeUI):
	# private enums
	__kDrawRectangle, __kDrawCircle, __kDrawTriangle = range(3)
	__kDrawWireframe, __kDrawWireframeOnShaded, __kDrawSmoothShaded, __kDrawFlatShaded, __kLastToken = range(5)

	def __init__(self):
		OpenMayaMPx.MPxSurfaceShapeUI.__init__(self)


	# override
	def getDrawRequests(self, info, objectAndActiveOnly, queue):
		"""
		 The draw data is used to pass geometry through the 
		 draw queue. The data should hold all the information
		 needed to draw the shape.
		"""
		data = OpenMayaUI.MDrawData()
		# printMsg("**before getProtoype\n");
		request = info.getPrototype(self)
		# printMsg("**after getProtoype\n");
		shapeNode = self.surfaceShape()
		geom = shapeNode.geometry()
		self.getDrawData(geom, data)
		request.setDrawData(data)

		# Are we displaying meshes?
		if (not info.objectDisplayStatus(OpenMayaUI.M3dView.kDisplayMeshes)):
			return

		# Use display status to determine what color to draw the object
		if (info.displayStyle() == OpenMayaUI.M3dView.kWireFrame):
			self.getDrawRequestsWireframe(request, info)
			queue.add(request)

		elif (info.displayStyle() == OpenMayaUI.M3dView.kGouraudShaded):
			request.setToken(basicShapeUI.__kDrawSmoothShaded)
			self.getDrawRequestsShaded(request, info, queue, data)
			queue.add(request)

		elif (info.displayStyle() == OpenMayaUI.M3dView.kFlatShaded):
			request.setToken(basicShapeUI.__kDrawFlatShaded)
			self.getDrawRequestsShaded(request, info, queue, data)
			queue.add(request)
		return 


	# override
	def draw(self, request, view):
		"""
		 From the given draw request, get the draw data and determine
		 which basic to draw and with what values.
		"""
		
		data = request.drawData()
		shapeNode = self.surfaceShape()
		geom = shapeNode.geometry()
		token = request.token()
		drawTexture = False

		su = OpenMaya.MScriptUtil()
		su.createFromDouble(1.0,1.0,1.0,1.0)

		#set up texturing if it is shaded
		if ((token == basicShapeUI.__kDrawSmoothShaded) or
					(token == basicShapeUI.__kDrawFlatShaded)):
			# Set up the material
			material = request.material()
			material.setMaterial(request.multiPath(), request.isTransparent())
				
			# Enable texturing
			# 
			# Note, Maya does not enable texturing when drawing with the 
			# default material. However, your custom shape is free to ignore
			# this setting.  
			# 
			drawTexture = material.materialIsTextured() and not view.usingDefaultMaterial() 
			
			# Apply the texture to the current view
			if (drawTexture):
				material.applyTexture(view, data)

		glFT.glPushAttrib( OpenMayaRender.MGL_ALL_ATTRIB_BITS )
		
		if ((token == basicShapeUI.__kDrawSmoothShaded) or
					(token == basicShapeUI.__kDrawFlatShaded)):			
			glFT.glEnable(OpenMayaRender.MGL_POLYGON_OFFSET_FILL)
			glFT.glPolygonMode(OpenMayaRender.MGL_FRONT_AND_BACK, OpenMayaRender.MGL_FILL)
			if (drawTexture):
				glFT.glEnable(OpenMayaRender.MGL_TEXTURE_2D)			
		else:
			glFT.glPolygonMode(OpenMayaRender.MGL_FRONT_AND_BACK, OpenMayaRender.MGL_LINE)
			
		# draw the shapes
		if (geom.shapeType == basicShapeUI.__kDrawCircle):
			# circle
			glFT.glBegin(OpenMayaRender.MGL_POLYGON)
			for i in range(0,360):
				rad = (i*2*math.pi)/360;
				glFT.glNormal3f(0.0, 0.0, 1.0)
				if (i == 360):
					glFT.glTexCoord3f(geom.radius*math.cos(0), geom.radius*math.sin(0), 0.0)
					glFT.glVertex3f(geom.radius*math.cos(0), geom.radius*math.sin(0), 0.0)
				else:
					glFT.glTexCoord3f(geom.radius*math.cos(rad), geom.radius*math.sin(rad), 0.0)
					glFT.glVertex3f(geom.radius*math.cos(rad), geom.radius*math.sin(rad), 0.0)
			glFT.glEnd()

		elif (geom.shapeType == basicShapeUI.__kDrawRectangle):
			#rectangle
			glFT.glBegin(OpenMayaRender.MGL_QUADS)

			glFT.glTexCoord2f(-1*(geom.width/2), -1*(geom.height/2))
			glFT.glVertex3f(-1*(geom.width/2), -1*(geom.height/2), 0.0)
			glFT.glNormal3f(0, 0, 1.0)
			
			glFT.glTexCoord2f(-1*(geom.width/2), (geom.height/2))
			glFT.glVertex3f(-1*(geom.width/2), (geom.height/2), 0.0)
			glFT.glNormal3f(0, 0, 1.0)
			
			glFT.glTexCoord2f((geom.width/2), (geom.height/2))
			glFT.glVertex3f((geom.width/2), (geom.height/2), 0.0)
			glFT.glNormal3f(0, 0, 1.0)
			
			glFT.glTexCoord2f((geom.width/2), -1*(geom.height/2))
			glFT.glVertex3f((geom.width/2), -1*(geom.height/2), 0.0)
			glFT.glNormal3f(0, 0, 1.0)
			glFT.glEnd()
		
		else:
			# triangle
			glFT.glBegin(OpenMayaRender.MGL_TRIANGLES)
			glFT.glTexCoord2f(-1*(geom.width/2), -1*(geom.height/2))
			glFT.glVertex3f(-1*(geom.width/2), -1*(geom.height/2), 0.0)
			glFT.glNormal3f(0.0, 0.0, 1.0)
			
			glFT.glTexCoord2f(0.0, (geom.height/2))
			glFT.glVertex3f(0.0, (geom.height/2), 0.0)
			glFT.glNormal3f(0.0, 0.0, 1.0)
			
			glFT.glTexCoord2f((geom.width/2), -1*(geom.height/2))
			glFT.glVertex3f((geom.width/2), -1*(geom.height/2), 0.0)
			glFT.glNormal3f(0.0, 0.0, 1.0)
			glFT.glEnd()
			
		if ((token == basicShapeUI.__kDrawSmoothShaded) or
					(token == basicShapeUI.__kDrawFlatShaded)):
			glFT.glDisable(OpenMayaRender.MGL_POLYGON_OFFSET_FILL)
			# Turn off texture mode
			if (drawTexture):
				glFT.glDisable(OpenMayaRender.MGL_TEXTURE_2D)
		
		glFT.glPopAttrib()



	# override
	def select(self, selectInfo, selectionList, worldSpaceSelectPts):
		"""
		 Select function. Gets called when the bbox for the object is selected.
		 This function just selects the object without doing any intersection tests.
		"""
		
		priorityMask = OpenMaya.MSelectionMask(OpenMaya.MSelectionMask.kSelectObjectsMask)
		item = OpenMaya.MSelectionList()
		item.add(selectInfo.selectPath())
		xformedPt = OpenMaya.MPoint()
		selectInfo.addSelection(item, xformedPt, selectionList,
								 worldSpaceSelectPts, priorityMask, False)
		return True


	def getDrawRequestsWireframe(self, request, info):
	
		request.setToken(basicShapeUI.__kDrawWireframe)

		displayStatus = info.displayStatus()
		activeColorTable = OpenMayaUI.M3dView.kActiveColors
		dormantColorTable = OpenMayaUI.M3dView.kDormantColors
		
		if (displayStatus == OpenMayaUI.M3dView.kLead):
				request.setColor(kLeadColor, activeColorTable)

		elif (displayStatus == OpenMayaUI.M3dView.kActive):
				request.setColor(kActiveColor, activeColorTable)

		elif (displayStatus == OpenMayaUI.M3dView.kActiveAffected):
				request.setColor(kActiveAffectedColor, activeColorTable)

		elif (displayStatus == OpenMayaUI.M3dView.kDormant):
				request.setColor(kDormantColor, dormantColorTable)

		elif (displayStatus == OpenMayaUI.M3dView.kHilite):
				request.setColor(kHiliteColor, activeColorTable)



	def getDrawRequestsShaded(self, request, info, queue, data):
		# Need to get the material info
		path = info.multiPath()	# path to your dag object 
		view = info.view() 		# view to draw to
		material = OpenMayaMPx.MPxSurfaceShapeUI.material(self, path)
		usingDefaultMat = view.usingDefaultMaterial()
		if usingDefaultMat: 
			material = OpenMayaUI.MMaterial.defaultMaterial() 

		displayStatus = info.displayStatus()

		# Evaluate the material and if necessary, the texture.
		try:
			material.evaluateMaterial(view, path)
		except RuntimeError:
			print "Couldn't evaluate material"
			raise

		drawTexture = not usingDefaultMat
		if (drawTexture and material.materialIsTextured()):
			material.evaluateTexture(data)

		request.setMaterial(material)

		#materialTransparent = False
		#material.getHasTransparency(materialTransparent)
		#if (materialTransparent):
		#	request.setIsTransparent(True)

		# create a draw request for wireframe on shaded if necessary.
		if ((displayStatus == OpenMayaUI.M3dView.kActive) or
				 (displayStatus == OpenMayaUI.M3dView.kLead) or
				 (displayStatus == OpenMayaUI.M3dView.kHilite)):
			wireRequest = info.getPrototype(self)
			wireRequest.setDrawData(data)
			self.getDrawRequestsWireframe(wireRequest, info)
			wireRequest.setToken(basicShapeUI.__kDrawWireframeOnShaded)
			wireRequest.setDisplayStyle(OpenMayaUI.M3dView.kWireFrame)
			queue.add(wireRequest)

#####################################################################

def nodeCreator():
	return OpenMayaMPx.asMPxPtr( basicShape() )


def uiCreator():
	return OpenMayaMPx.asMPxPtr( basicShapeUI() )


def nodeInitializer():
	# BASIC type enumerated attribute
	enumAttr = OpenMaya.MFnEnumAttribute()
	basicShape.aShapeType = enumAttr.create("shapeType", "st", kDefaultShapeType)
	enumAttr.addField("rectangle", 0)
	enumAttr.addField("circle", 1)
	enumAttr.addField("triangle", 2)
	enumAttr.setHidden(False)
	enumAttr.setKeyable(True)
	basicShape.addAttribute(basicShape.aShapeType)

	# BASIC numeric attributes
	# utility func for numeric attrs
	def setOptions(attr):
		attr.setHidden(False)
		attr.setKeyable(True)
		attr.setInternal(True)

	numericAttr = OpenMaya.MFnNumericAttribute()

	basicShape.aRadius = numericAttr.create("radius", "r", OpenMaya.MFnNumericData.kDouble, kDefaultRadius)
	setOptions(numericAttr)
	basicShape.addAttribute(basicShape.aRadius)

	basicShape.aHeight = numericAttr.create("height", "ht", OpenMaya.MFnNumericData.kDouble, kDefaultHeight)
	setOptions(numericAttr)
	basicShape.addAttribute(basicShape.aHeight)
	
	basicShape.aWidth = numericAttr.create("width2", "wt2", OpenMaya.MFnNumericData.kDouble, kDefaultWidth)
	setOptions(numericAttr)
	basicShape.addAttribute(basicShape.aWidth)

# initialize the script plug-in
def initializePlugin(mobject):
	mplugin = OpenMayaMPx.MFnPlugin(mobject, "Autodesk", "8.5", "Any")
	try:
		mplugin.registerShape( kPluginNodeTypeName, spBasicShapeNodeId,
								nodeCreator, nodeInitializer, uiCreator )
	except:
		sys.stderr.write( "Failed to register node: %s" % kPluginNodeTypeName )
		raise


# uninitialize the script plug-in
def uninitializePlugin(mobject):
	mplugin = OpenMayaMPx.MFnPlugin(mobject)
	try:
		mplugin.deregisterNode( spBasicShapeNodeId )
	except:
		sys.stderr.write( "Failed to deregister node: %s" % kPluginNodeTypeName )
		raise
